# Copyright 2025 OpenAccess AI Collective and the LlamaFactory team.
#
# This code is inspired by the OpenAccess AI Collective's axolotl library.
# https://github.com/OpenAccess-AI-Collective/axolotl/blob/main/src/axolotl/monkeypatch/utils.py
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Literal, Optional

import numpy as np
import torch
import torch.nn.functional as F
from transformers import DataCollatorForSeq2Seq

from ..extras.constants import AUDIO_PLACEHOLDER, IGNORE_INDEX, IMAGE_PLACEHOLDER
from ..extras.packages import is_pillow_available


if is_pillow_available():
    from PIL import Image


if TYPE_CHECKING:
    from transformers import ProcessorMixin

    from .template import Template


def prepare_4d_attention_mask(attention_mask_with_indices: "torch.Tensor", dtype: "torch.dtype") -> "torch.Tensor":
    r"""Expand 2d attention mask to 4d attention mask.

    Expand the attention mask with indices from (batch_size, seq_len) to (batch_size, 1, seq_len, seq_len),
    handle packed sequences and transforms the mask to lower triangular form to prevent future peeking.

    e.g.
    ```python
    # input
    [[1, 1, 2, 2, 2, 0]]
    # output
    [
        [
            [
                [o, x, x, x, x, x],
                [o, o, x, x, x, x],
                [x, x, o, x, x, x],
                [x, x, o, o, x, x],
                [x, x, o, o, o, x],
                [x, x, x, x, x, x],
            ]
        ]
    ]
    ```
    where `o` equals to `0.0`, `x` equals to `min_dtype`.
    """
    _, seq_len = attention_mask_with_indices.size()
    min_dtype = torch.finfo(dtype).min
    zero_tensor = torch.tensor(0, dtype=dtype)

    # Create a non-padding mask.
    non_padding_mask = (attention_mask_with_indices != 0).unsqueeze(1).unsqueeze(2)
    # Create indices for comparison.
    indices = attention_mask_with_indices.unsqueeze(1).unsqueeze(2)  # [bsz, 1, 1, seq_len]
    indices_t = attention_mask_with_indices.unsqueeze(1).unsqueeze(3)  # [bsz, 1, seq_len, 1]
    # Create a lower triangular mask.
    tril_mask = torch.tril(torch.ones((seq_len, seq_len), dtype=torch.bool))
    attention_mask_4d = (indices == indices_t) & non_padding_mask & tril_mask
    # Invert the attention mask.
    attention_mask_4d = torch.where(attention_mask_4d, zero_tensor, min_dtype)
    return attention_mask_4d


@dataclass
class MultiModalDataCollatorForSeq2Seq(DataCollatorForSeq2Seq):
    r"""Data collator that supports VLMs.

    Features should contain input_ids, attention_mask, labels, and optionally contain images, videos and audios.
    """

    # 可选的模板对象，用于多模态输入处理
    template: Optional["Template"] = None
    # 可选的处理器对象（如 HuggingFace 的 Processor），用于图像/音频等预处理
    processor: Optional["ProcessorMixin"] = None

    # 后初始化 
    def __post_init__(self):
        """
        初始化后检查是否提供了模板。
        如果没有提供模板，则抛出异常，因为它是处理多模态输入所必需的。
        """
        # 如果未提供 template，则抛出 ValueError 异常
        if self.template is None:
            raise ValueError("Template is required for MultiModalDataCollator.")

    def __call__(self, features: list[dict[str, Any]]) -> dict[str, "torch.Tensor"]:
        # 存储所有样本中的图像、视频和音频
        batch_images, batch_videos, batch_audios = [], [], []
        # 记录每种类型的数据长度
        batch_imglens, batch_vidlens, batch_audlens, batch_input_ids = [], [], [], []
        # 遍历每个样本，提取并记录多模态数据
        for feature in features:
            # 提取并删除 feature 中的 images 字段，若不存在则返回空列表
            images = feature.pop("images", None) or []
            # 提取并删除 videos 字段
            videos = feature.pop("videos", None) or []
            # 提取并删除 audios 字段
            audios = feature.pop("audios", None) or []
            # 将当前样本的图像添加到总 batch 图像列表中
            batch_images.extend(images)
            # 添加视频
            batch_videos.extend(videos)
            # 添加音频
            batch_audios.extend(audios)
            # 记录当前样本包含的图像数量
            batch_imglens.append(len(images))
            # 视频数量
            batch_vidlens.append(len(videos))
            # 音频数量
            batch_audlens.append(len(audios))
            # 记录 input_ids
            batch_input_ids.append(feature["input_ids"])

        # 如果没有图像或视频但需要时，添加假图像避免进程挂起
        fake_input_ids = []
        # 如果模型需要 image_token 但当前没有图像或视频，生成假图像避免 Zero3/FSDP 挂起
        if (
            self.template.mm_plugin.image_token is not None and sum(batch_imglens) == 0 and sum(batch_vidlens) == 0
        ):  # avoid process hanging in zero3/fsdp case
            # 创建一个假消息和空白图像作为占位符
            fake_messages = [{"role": "user", "content": IMAGE_PLACEHOLDER}]
            fake_images = [Image.new("RGB", (64, 64), (255, 255, 255))]
            # 使用 mm_plugin 处理假消息，加入图像信息
            fake_messages = self.template.mm_plugin.process_messages(
                fake_messages, fake_images, [], [], self.processor
            )
            # 对假消息内容进行编码，得到假 input_ids
            _fake_input_ids = self.tokenizer.encode(fake_messages[0]["content"], add_special_tokens=False)
            # 继续用 mm_plugin 处理 token IDs 和图像
            _fake_input_ids, _ = self.template.mm_plugin.process_token_ids(
                _fake_input_ids, None, fake_images, [], [], self.tokenizer, self.processor
            )
            # 将假 input_ids 添加到 fake_input_ids 列表中
            fake_input_ids.extend(_fake_input_ids)
            # 替换当前 batch 中的图像为假图像
            batch_images = fake_images
            # 更新图像长度
            batch_imglens[0] = 1

        # 类似地，如果没有音频但需要时，添加假音频避免进程挂起
        # 如果模型需要 audio_token 但当前没有音频，生成假音频防止挂起
        if (
            self.template.mm_plugin.audio_token is not None and sum(batch_audlens) == 0
        ):  # avoid process hanging in zero3/fsdp case
            # 构造假消息和零值音频数组
            fake_messages = [{"role": "user", "content": AUDIO_PLACEHOLDER}]
            fake_audios = [np.zeros(1600)]
            # 使用 mm_plugin 处理假消息和音频
            fake_messages = self.template.mm_plugin.process_messages(
                fake_messages, [], [], fake_audios, self.processor
            )
            # 编码假消息内容得到 input_ids
            _fake_input_ids = self.tokenizer.encode(fake_messages[0]["content"], add_special_tokens=False)
            # 使用 mm_plugin 处理 token IDs 和音频
            _fake_input_ids, _ = self.template.mm_plugin.process_token_ids(
                _fake_input_ids, None, [], [], fake_audios, self.tokenizer, self.processor
            )
            # 添加假 input_ids
            fake_input_ids.extend(_fake_input_ids)
            # 替换音频为假音频
            batch_audios = fake_audios
            # 更新音频长度
            batch_audlens[0] = 1

        # 如果有假输入ID，则根据tokenizer的padding方向将其插入到第一个样本中
        if len(fake_input_ids) != 0:
            # 若 padding 在右边，则将假 input_ids 加在 input_ids 后面
            if self.tokenizer.padding_side == "right":
                features[0]["input_ids"] = features[0]["input_ids"] + fake_input_ids
                features[0]["attention_mask"] = features[0]["attention_mask"] + [0] * len(fake_input_ids)
                features[0]["labels"] = features[0]["labels"] + [IGNORE_INDEX] * len(fake_input_ids)
            else:
                # padding 在左边，则插入前面
                features[0]["input_ids"] = fake_input_ids + features[0]["input_ids"]
                features[0]["attention_mask"] = [0] * len(fake_input_ids) + features[0]["attention_mask"]
                features[0]["labels"] = [IGNORE_INDEX] * len(fake_input_ids) + features[0]["labels"]

            # 更新 batch_input_ids
            batch_input_ids[0] = features[0]["input_ids"]

        # 调用 mm_plugin 获取最终的多模态输入张量
        mm_inputs = self.template.mm_plugin.get_mm_inputs(
            batch_images,
            batch_videos,
            batch_audios,
            batch_imglens,
            batch_vidlens,
            batch_audlens,
            batch_input_ids,
            self.processor,
        )
        # 如果有 token_type_ids，将其单独取出，并插入到每个 feature 中
        if "token_type_ids" in mm_inputs:
            token_type_ids = mm_inputs.pop("token_type_ids")
            for i, feature in enumerate(features):
                feature["token_type_ids"] = token_type_ids[i]

        # 调用父类的数据整理逻辑，处理 input_ids、labels、attention_mask 等
        features: dict[str, torch.Tensor] = super().__call__(features)

        # 如果模型支持 get_rope_index（如 qwen2vl），计算 RoPE 相关索引
        if self.model is not None and hasattr(self.model, "get_rope_index"):  # for qwen2vl mrope
            # 构建参数字典，用于获取 rope_index
            rope_index_kwargs = {
                "input_ids": features["input_ids"],
                "image_grid_thw": mm_inputs.get("image_grid_thw"),
                "video_grid_thw": mm_inputs.get("video_grid_thw"),
                "attention_mask": (features["attention_mask"] >= 1).float(),
            }
            # 如果有 second_per_grid_ts，加入 kwargs
            if "second_per_grid_ts" in mm_inputs:  # for qwen2vl
                rope_index_kwargs["second_per_grid_ts"] = mm_inputs.get("second_per_grid_ts")
            # 如果有 video_second_per_grid，也加入
            if "video_second_per_grid" in mm_inputs:  # for qwen2omni
                rope_index_kwargs["second_per_grids"] = mm_inputs.get("video_second_per_grid")

            # 特别处理 qwen2_5_omni_thinker 类型的模型
            if getattr(self.model.config, "model_type", None) == "qwen2_5_omni_thinker":  # for qwen2omni
                # 设置 use_audio_in_video 参数
                rope_index_kwargs["use_audio_in_video"] = getattr(self.processor, "use_audio_in_video", False)
                # 获取特征注意力掩码
                feature_attention_mask = mm_inputs.get("feature_attention_mask", None)
                if feature_attention_mask is not None:
                    # 计算音频特征长度
                    audio_feature_lengths = torch.sum(
                        feature_attention_mask, dim=1
                    )  # FIXME need to get video image lengths
                    # 加入音频序列长度
                    rope_index_kwargs["audio_seqlens"] = audio_feature_lengths  # prepare for input

                # 计算 delta 值以避免冲突
                delta0 = (1 - rope_index_kwargs["attention_mask"]).sum(dim=-1).unsqueeze(1)
                # avoid conflict
                # 获取 position_ids 和 rope_deltas
                new_position_ids, rope_deltas = self.model.get_rope_index(**rope_index_kwargs)
                features["position_ids"], features["rope_deltas"] = (
                    new_position_ids.clone(),
                    rope_deltas - delta0,
                )  # avoid inplace operation FIXME
            else:  # for qwen2vl
                # 正常调用 get_rope_index 获取 position_ids 和 rope_deltas
                features["position_ids"], features["rope_deltas"] = self.model.get_rope_index(**rope_index_kwargs)

        # 如果有 cross_attention_mask，做 pad 操作以适配 LLaMA 的输入需求
        if "cross_attention_mask" in mm_inputs:  # for mllama inputs when pad_to_multiple_of is enabled
            cross_attention_mask = mm_inputs.pop("cross_attention_mask")
            seq_len = features["input_ids"].size(1)
            orig_len = cross_attention_mask.size(1)
            mm_inputs["cross_attention_mask"] = F.pad(cross_attention_mask, (0, 0, 0, 0, 0, seq_len - orig_len))

        # 将 mm_inputs 合并进 features 字典中
        features.update(mm_inputs)

        # 如果有 image_bound（如 minicpmv 模型），重新构造 position_ids
        if "image_bound" in features:  # for minicpmv inputs
            bsz, seq_length = features["input_ids"].shape
            features["position_ids"] = torch.arange(seq_length).long().repeat(bsz, 1)
            return {"data": features, "input_ids": features["input_ids"], "labels": features["labels"]}

        # 返回最终的 batch 数据字典
        return features


@dataclass
class SFTDataCollatorWith4DAttentionMask(MultiModalDataCollatorForSeq2Seq):
    r"""Data collator for 4d attention mask."""
    # 定义一个支持 4D attention mask 的数据整理器（Data Collator），用于 SFT 训练

    # 是否使用 block diagonal attention（块对角注意力），用于多模态或多任务场景
    block_diag_attn: bool = False
    # 指定使用的注意力实现方式：eager、sdpa 或 flash_attention_2
    attn_implementation: Literal["eager", "sdpa", "flash_attention_2"] = "eager"
    # 指定计算时使用的数据类型，默认是 torch.float32
    compute_dtype: "torch.dtype" = torch.float32

    def __call__(self, features: list[dict[str, Any]]) -> dict[str, "torch.Tensor"]:
        # 调用父类的数据整理逻辑，处理输入特征（如 padding、tensor 化等）
        features = super().__call__(features)
        # 如果启用 block_diag_attn 且当前不是使用 flash_attention_2，则构建 4D attention mask
        if self.block_diag_attn and self.attn_implementation != "flash_attention_2":
            features["attention_mask"] = prepare_4d_attention_mask(features["attention_mask"], self.compute_dtype)

        # 遍历所有特征张量，将浮点型张量转换为指定的计算精度（如 float16 或 bfloat16）
        for key, value in features.items():  # cast data dtype for paligemma
            if torch.is_tensor(value) and torch.is_floating_point(value):
                features[key] = value.to(self.compute_dtype)

        # 返回整理后的 batch 数据
        return features


@dataclass
class PairwiseDataCollatorWithPadding(MultiModalDataCollatorForSeq2Seq):
    r"""Data collator for pairwise data."""

    def __call__(self, features: list[dict[str, Any]]) -> dict[str, "torch.Tensor"]:
        r"""Pad batched data to the longest sequence in the batch.

        We generate 2 * n examples where the first n examples represent chosen examples and
        the last n examples represent rejected examples.
        """
        concatenated_features = []
        for key in ("chosen", "rejected"):
            for feature in features:
                target_feature = {
                    "input_ids": feature[f"{key}_input_ids"],
                    "attention_mask": feature[f"{key}_attention_mask"],
                    "labels": feature[f"{key}_labels"],
                    "images": feature["images"],
                    "videos": feature["videos"],
                    "audios": feature["audios"],
                }
                concatenated_features.append(target_feature)

        return super().__call__(concatenated_features)


@dataclass
class KTODataCollatorWithPadding(MultiModalDataCollatorForSeq2Seq):
    r"""Data collator for KTO data."""

    def __call__(self, features: list[dict[str, Any]]) -> dict[str, "torch.Tensor"]:
        target_features = []
        kl_features = []
        kto_tags = []
        for feature in features:
            target_feature = {
                "input_ids": feature["input_ids"],
                "attention_mask": feature["attention_mask"],
                "labels": feature["labels"],
                "images": feature["images"],
                "videos": feature["videos"],
                "audios": feature["audios"],
            }
            kl_feature = {
                "input_ids": feature["kl_input_ids"],
                "attention_mask": feature["kl_attention_mask"],
                "labels": feature["kl_labels"],
                "images": feature["images"],
                "videos": feature["videos"],
                "audios": feature["audios"],
            }
            target_features.append(target_feature)
            kl_features.append(kl_feature)
            kto_tags.append(feature["kto_tags"])

        batch = super().__call__(target_features)
        kl_batch = super().__call__(kl_features)
        batch["kl_input_ids"] = kl_batch["input_ids"]
        batch["kl_attention_mask"] = kl_batch["attention_mask"]
        batch["kl_labels"] = kl_batch["labels"]
        if "cross_attention_mask" in kl_batch:  # for mllama inputs
            batch["kl_cross_attention_mask"] = kl_batch["cross_attention_mask"]

        if "token_type_ids" in kl_batch:
            batch["kl_token_type_ids"] = kl_batch["token_type_ids"]

        batch["kto_tags"] = torch.tensor(kto_tags)
        return batch
